#coding=utf-8

import numpy as np
from scipy.spatial.transform import Rotation
class GeoEntity:
    '''
    点
    '''
    def __init__(self,x=0,y=0,z=0,roll=0,pitch=0,yaw=0,cacheSize=5,name="J"):
        self.x=x
        self.y=y
        self.z=z
        self.roll=roll
        self.pitch=pitch
        self.yaw=yaw
        self.name=name
        self.posQueue=[]
        self.rotQueue=[]
        for i in range(cacheSize):
            self.posQueue.append(np.zeros(3))
            self.rotQueue.append(np.zeros(3))

    def DumpCache(self):
        self.posQueue.append(self.VectorPose())
        self.rotQueue.append(self.VectorRot())
        self.posQueue.pop(0)
        self.rotQueue.pop(0)

    def __sub__(self, other):
        return GeoEntity(other.x-self.x,other.y-self.y,other.z-self.z,other.roll-self.roll,other.pitch-self.pitch,other.yaw-self.yaw)
    def __add__(self, other):
        return GeoEntity(other.x + self.x, other.y + self.y, other.z + self.z,other.roll+self.roll,other.pitch+self.pitch,other.yaw+self.yaw)
    def Rotate(self,r,p,y):
        '''
        自转
        '''
        self.roll+=r;self.pitch+=p;self.yaw+=y
    def Translate(self,x,y,z):
        '''
        平移
        '''
        self.x+=x;self.y+=y;self.z+=z
    def Distance(self,entityB):
        '''
        两点之间的欧拉距离
        :return:
        '''
        return np.linalg.norm(np.array([self.x-entityB.x,self.y-entityB.y,self.z-entityB.z]))

    def DistancebyAxis(self,entity):
        '''
        逐个轴计算插值
        :return:
        '''
        return self.x-entity.x,self.y-entity.y,self.z-entity.z
    def MoveTo(self,target,distance):
        x,y,z=target.DistancebyAxis(self)
        total=target.Distance(self)
        if total==0: return
        dx=x*distance/total
        dy=y*distance/total
        dz=z*distance/total
        self.Translate(dx,dy,dz)

    def VectorPose(self):
        return np.array([self.x,self.y,self.z])
    def VectorRot(self):
        return np.array([self.roll,self.pitch,self.yaw])
    def Parallel(self,entity):
        '''
        调整姿态，与一个实体平行
        :param entity:
        :return:
        '''
        self.pitch=entity.pitch
        self.roll=entity.roll
        self.yaw=entity.yaw

    def SetGeoData(self,position=np.zeros(3),rotation=np.zeros(3),pos=True,rot=True):
        '''
        初始化几何学位置
        :param position:
        :param rotation:
        :return:
        '''
        if pos:
            self.SetPose(position)
        if rot:
            self.SetRot(rotation)
    def SetPose(self,position=np.zeros(3)):
        '''
        从ndarray或列表中更新位置信息
        :param position:
        :return:
        '''
        self.x = position[0]
        self.y = position[1]
        self.z = position[2]
    def SetRot(self,rotation=np.zeros(3)):
        '''
        从ndarray或列表中更新旋转信息
        :param rotation:
        :return:
        '''
        self.roll = rotation[0]
        self.pitch = rotation[1]
        self.yaw = rotation[2]
    def Mimic(self,entity):
        '''
        模仿控件属性
        :param entity:
        :return:
        '''
        self.SetGeoData(entity.VectorPose(),entity.VectorRot())
    def LookupTransform(self,entity):
        '''
        获取一个实体相对于另一个实体的空间变换
        :param entity: 参考实体
        :return: x,y,z,roll,pitch,yaw
        '''
        '''
            下面这一大堆还是有问题，等着研究一下
        '''
        thisRMatrix=EulerMatrix(self.roll,self.pitch,self.yaw)
        thatRMatrix=EulerMatrix(entity.roll,entity.pitch,entity.yaw)

        thisTransform=np.eye(4)
        thatTransform=np.eye(4)

        thisTransform[:3,:3]=thisRMatrix
        thisTransform[:3,3]=[self.x,self.y,self.z]
        thatTransform[:3,:3]=thatRMatrix
        thatTransform[:3,3]=[entity.x,entity.y,entity.z]

        relativeTrans=(np.linalg.inv(thisTransform))@thatTransform
        x,y,z=relativeTrans[:3,3]
        rotMatrix=relativeTrans[:3,:3]
        roll,pitch,yaw=RestoreRPYFromMatrix(rotMatrix)
        return x,y,z,roll,pitch,yaw
    def BornCenterEntity(self,entity):
        '''
        这个实体和另一个实体，生成中间实体
        :param entity:
        :return:
        '''
        e=GeoEntity()
        pos=(self.VectorPose()+entity.VectorPose())/2
        e.SetPose(pos)
        return e
    def __str__(self):
        '''显示信息'''
        return '实体坐标[{:.6f},{:.6f},{:.6f}],旋转[{:.6f},{:.6f},{:.6f}]'.format(self.x,self.y,self.z,self.roll,self.pitch,self.yaw)

class GeoLinks:
    '''
    两个GeoEntity之间的连接
    '''
    def __init__(self,startX=0,startY=0,startZ=0,length=0,name="J"):
        '''
        初始化一段link
        '''
        self.start=GeoEntity(startX,startY,startZ)
        self.end=GeoEntity(startX,startY,startZ)
        self.length=length
        self.name=name
        self.UpdateEnd()



    def SetEndGlobalRotate(self,localRoll,localPitch,localYaw,order='xyz'):
        '''
        设置末端旋转值
        :param roll:
        :param pitch:
        :param yaw:
        :return:
        '''
        self.end.roll,self.end.pitch,self.end.yaw=Local2GlobalRotate(-localRoll,-localPitch,-localYaw,self.start.roll,self.start.pitch,self.start.yaw,order)
    def SetStartGlobalRotate(self,localRoll,localPitch,localYaw,order='xyz'):
        '''
        设置末端旋转值
        :param roll:
        :param pitch:
        :param yaw:
        :return:
        '''
        self.start.roll,self.start.pitch,self.start.yaw=Local2GlobalRotate(localRoll, localPitch, localYaw, self.end.roll, self.end.pitch, self.end.yaw,
                                  order)
    def UpdateEnd(self):
        '''
        更新末端
        :return:
        '''

        rotated_point = Rotate3D(self.start.roll,self.start.pitch,self.start.yaw,0,0,self.length)
        self.end.SetPose(rotated_point + self.start.VectorPose())
        self.end.SetRot(self.start.VectorRot())
    def UpdateStart(self):
        '''
        基于末端更新前端
        :param pitch:
        :param yaw:
        :return:
        '''
        rotated_point = Rotate3D(self.start.roll,self.start.pitch,self.start.yaw,0,0,-self.length)
        self.start.SetPose(rotated_point+ self.end.VectorPose())

    def AppendTo(self,previousLink):
        '''
        追加到另一个Link之上
        :param previousLink:
        :return:
        '''
        self.start.Mimic(previousLink.end)
        self.UpdateEnd()

    def Vecotorize(self):
        '''
        类向量化[2x3矩阵]
        [[起点],[终点]]
        :return:
        '''
        return np.array([self.start.VectorPose(),self.end.VectorPose()])
    def CenterPose(self):
        '''
        返回link的中心点
        :return:
        '''
        return (self.start.VectorPose()+self.end.VectorPose())/2

def RestoreRPYFromMatrix(matrix):
    '''
    从旋转矩阵中解析出Roll,Pitch,Yaw
    :param matrix:
    :return:
    '''
    roll = np.arctan2(matrix[2, 1], matrix[2, 2])
    pitch = np.arctan2(-matrix[2, 0], np.sqrt(matrix[2, 1] ** 2 + matrix[2, 2] ** 2))
    yaw = np.arctan2(matrix[1, 0], matrix[0, 0])
    return roll,pitch,yaw
def calculate_tilt_angles(aX, aY, aZ):
    # 计算合加速度

    tiltX=np.arctan2(aY,aZ)
    tiltY=np.arctan2(-aX,np.sqrt(aY**2 + aZ**2))

    return np.array([tiltX, tiltY, 0])


def Rotate3D(roll=0, pitch=0, yaw=0,x=0,y=0,z=0,order='xyz'):
    '''
    在Entity上面放置一个长度为l的bone，然后旋转骨骼，接着沿骨骼点转移，得到下一个Entity的位置和姿态
    :param l:   骨骼长度
    :param roll:
    :param pitch:
    :param yaw:
    :return:
    '''
    initial_point = np.array([x, y, z]).transpose()
    mat=EulerMatrix(roll,pitch,yaw,order)
    p=mat@initial_point
    return p.transpose()

def calculate_triangle_angles(a, b, c):
    # 计算余弦值
    cos_A = (b**2 + c**2 - a**2) / (2 * b * c)
    cos_B = (a**2 + c**2 - b**2) / (2 * a * c)
    cos_C = (a**2 + b**2 - c**2) / (2 * a * b)

    # 通过反余弦函数计算角度（注意：结果为弧度）
    angle_A = np.arccos(cos_A)
    angle_B = np.arccos(cos_B)
    angle_C = np.arccos(cos_C)

    return angle_A, angle_B, angle_C

def EulerMatrix(roll, pitch, yaw,order='xyz'):
    # 将欧拉角转换为旋转矩阵
    matrixBundle = {}
    #   yaw
    matrixBundle['z'] = np.array([
        [np.cos(yaw), -np.sin(yaw), 0],
        [np.sin(yaw), np.cos(yaw), 0],
        [0, 0, 1]
    ])
    #   pitch
    matrixBundle['y'] = np.array([
        [np.cos(pitch), 0, np.sin(pitch)],
        [0, 1, 0],
        [-np.sin(pitch), 0, np.cos(pitch)]
    ])
    #   roll
    matrixBundle['x'] = np.array([
        [1, 0, 0],
        [0, np.cos(roll), -np.sin(roll)],
        [0, np.sin(roll), np.cos(roll)]
    ])
    mat = np.eye(3)
    for axis in order:
        mat = matrixBundle[axis] @ mat
    return mat
def Local2GlobalRotate(roll,pitch,yaw,groll,gpitch,gyaw,order='xyz'):
    '''
    设置末端旋转值
    :param roll:
    :param pitch:
    :param yaw:
    :return:
    '''
    localMatrix=EulerMatrix(roll,pitch,yaw,order)
    local2GlobalMatrix=EulerMatrix(groll,gpitch,gyaw,order)
    globalMatrix=local2GlobalMatrix@localMatrix
    roll,pitch,yaw=RestoreRPYFromMatrix(globalMatrix)
    return roll,pitch,yaw
def InRange(x,y,xRange,yRange):
    '''
    判断X，Y是否在范围内
    :param x:
    :param y:
    :param xRange:
    :param yRange:
    :return:
    '''
    if x>xRange[0] and x<xRange[1] and y>yRange[0] and y<yRange[1]:
        return True
    return False

if __name__=='__main__':
    '''
    debug
    0.0 0.01738249977583506 -0.2521913370029726
    '''
    # a=GeoEntity(0,0,0.0,np.pi/2,0,0.0)
    # x=GeoLinks(length=1)
    # x.end=a
    # x.SetStartGlobalRotate(0,np.pi/2,0)
    # x.UpdateStart()
    # print(x.start)
    roll=np.pi/4
    pitch=roll
    yaw=0
    #for yaw in np.linspace(0,np.pi,10):
    m=EulerMatrix(roll,pitch,yaw)
    print(RestoreRPYFromMatrix(m))